import os
import uuid
from typing import Any, Optional

from configs import dify_config
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType
from core.model_runtime.errors.invoke import (
    InvokeAuthorizationError,
    InvokeBadRequestError,
    InvokeConnectionError,
    InvokeError,
    InvokeRateLimitError,
    InvokeServerUnavailableError,
)
from core.model_runtime.model_providers.__base.tts_model import TTSModel
from core.model_runtime.model_providers.picc.helper.picc_model_center import TTSClient


class PiccText2SpeechModel(TTSModel):
    voiceIdMap = {
        "yifei": "62320",
        "chunchun": "65580",
        "jingjing": "62360",
        "panting": "62420",
        "xiaotao": "67910",
        "xiaoruan": "62340",
        "xiaoyi": "65600",
    }

    def _invoke(
        self, model: str, tenant_id: str, credentials: dict, content_text: str, voice: str, user: Optional[str] = None
    ) -> Any:
        """
        _invoke text2speech model

        :param model: model name
        :param tenant_id: user tenant id
        :param credentials: model credentials
        :param content_text: text content to be translated
        :param voice: model timbre
        :param user: unique user id
        :return: text translated to audio file
        """
        return self._tts_invoke_streaming(model=model, credentials=credentials, content_text=content_text, voice=voice)

    def validate_credentials(self, model: str, credentials: dict, user: Optional[str] = None) -> None:
        """
        validate credentials text2speech model

        :param model: model name
        :param credentials: model credentials
        :param user: unique user id
        :return: text translated to audio file
        """
        try:
            url = credentials["base_url"]
            AbilityId = credentials["ability_id"]
            AppId = credentials["app_id"]
            SecretId = credentials["secret_id"]
            SecretKey = credentials["secret_key"]
            client = TTSClient(url, AbilityId, AppId, SecretId, SecretKey)
            text = "测试连接"
            bytes = client.invoke(text=text)
            if bytes == b'':
                raise InvokeAuthorizationError("认证信息错误")
        except Exception as ex:
            raise InvokeAuthorizationError(str(ex))

    def _tts_invoke_streaming(self, model: str, credentials: dict, content_text: str, voice: str) -> Any:
        """
        _tts_invoke_streaming text2speech model

        :param model: model name
        :param credentials: model credentials
        :param content_text: text content to be translated
        :param voice: model timbre
        :return: text translated to audio file
        """
        url = credentials["base_url"]
        AbilityId = credentials["ability_id"]
        AppId = credentials["app_id"]
        SecretId = credentials["secret_id"]
        SecretKey = credentials["secret_key"]
        client = TTSClient(url, AbilityId, AppId, SecretId, SecretKey)
        voiceId = self.voiceIdMap[voice]
        tempFilePath = dify_config.PICC_FILE_TEMP_PATH + "temp/"
        if not os.path.exists(tempFilePath):
            os.mkdir(tempFilePath)
        tempFile = tempFilePath + str(uuid.uuid4()) + ".wav"
        # data = client.invoke(text)
        success = client.invoke(text=content_text, to_file_path=tempFile, voiceId=voiceId, volume=10, speed=-500, pitch=-500)

        data = []
        if success:
            with open(tempFile, "rb") as file:
                data = file.read()
        for i in range(0, len(data), 1024):
            yield data[i : i + 1024]
        os.remove(tempFile)

    def _process_sentence(self, sentence: str, model: str, voice, credentials: dict):
        """
        _tts_invoke picc text2speech model api

        :param model: model name
        :param credentials: model credentials
        :param voice: model timbre
        :param sentence: text content to be translated
        :return: text translated to audio file
        """
        pass

    def get_customizable_model_schema(self, model: str, credentials: dict) -> Optional[AIModelEntity]:
        """
        Get customizable model schema
        """
        model_schema = self.get_model_schema("TTS(picc)", credentials)
        if not model_schema or ModelPropertyKey.VOICES not in model_schema.model_properties:
            raise ValueError("this model does not support voice")

        voices = model_schema.model_properties[ModelPropertyKey.VOICES]  # If no voices provided or all voices were empty strings, use 'alloy' as default
        if not voices:
            voices = [{"name": "Alloy", "mode": "alloy", "language": "en-US"}]
        return AIModelEntity(
            model=model,
            label=I18nObject(zh_Hans=model, en_US=model),
            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
            model_type=ModelType.TTS,
            model_properties={
                ModelPropertyKey.AUDIO_TYPE: "wav",
                ModelPropertyKey.WORD_LIMIT: 4096,
                ModelPropertyKey.DEFAULT_VOICE: voices[0]["mode"],
                ModelPropertyKey.VOICES: voices,
            },
        )

    @property
    def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
        """
        Map model invoke error to unified error
        The key is the error type thrown to the caller
        The value is the error type thrown by the model,
        which needs to be converted into a unified error type for the caller.

        :return: Invoke error mapping
        """
        return {
            InvokeConnectionError: [InvokeConnectionError],
            InvokeServerUnavailableError: [InvokeServerUnavailableError],
            InvokeRateLimitError: [InvokeRateLimitError],
            InvokeAuthorizationError: [InvokeAuthorizationError],
            InvokeBadRequestError: [InvokeBadRequestError, KeyError, ValueError],
        }





